program program_psr
use madx_ptc_module
use pointer_lattice
implicit none

    interface
       subroutine build_tiny_package_example(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_tiny_package_example
    end interface

    interface
       subroutine build_PSR(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_PSR
    end interface

type(layout), pointer:: psr
type(damap) id,m,a_t
type(real_8) ray(6)

type(internal_state) state
type(fibre), pointer :: f
type(normalform) normal
integer :: np1=0,je(lnv),i,j,k,mf
integer,parameter :: no1=3 ! The order us hardwired in this example
real(dp) closed_orbit(6),disp(4,no1)

logical(lp) mis

call ptc_ini

psr=>m_u%start

call kanalnummer(mf,"Dispersions.txt")
Write(6,*) " small misalignments and rotations in lattice ? input t or f "
read(5,*) mis
 
 
 call build_PSR(PSR,mis)
 

!!!  Phase space of Taylor maps will be (x,px,y,py,delta)
state=delta0

j=0
np1=0    !!!  No system parameters in the map

call init(state,no1,np1) ! PTC and FPP are properly initialized

call alloc(id,m,A_T)  
call alloc(ray)
call alloc(normal)


closed_orbit=0.d0 ! initial guess for closed orbit

call find_orbit(psr,closed_orbit,1,state,1.d-5)

write(6,'(a16,4(1x,g12.5))') " closed orbit = ",closed_orbit(1:4)

id=1;   ray=closed_orbit+id;   ! ray= closed orbit + identity map

call track(psr,ray,1,state) ! One turn map is computed via the ray

!!! The ray is really the result of a truncated power series algebra
!!! The map we are interested to compute is around the closed orbit: it is really part of a diffential algebra
m=ray   !   The ray is "officially" turned into a "damap"  (DA in Berz's honour)

normal = m  ! The map is normalised   


 
 !  Here starts a universal Twiss loop
 f=>psr%start               

 ray=normal%a_t+closed_orbit   ! The ray is initialize as the map A_cs+closed_orbit


  a_t=ray

 do j=1,no1
  je=0
  je(5)=j
  do k=1,4
   disp(k,j)=a_t%v(k).sub.je  ! All the dispersions up to delta^3 are stored in disp(1:4,1:3)
  enddo
enddo
  write(mf,*) " At the beginning of the lattice " 
  write(mf,'(a19,3(1x,D18.11))') " x  dispersions -> ",disp(1,:)
  write(mf,'(a19,3(1x,D18.11))') " x' dispersions -> ",disp(2,:)
  write(mf,'(a19,3(1x,D18.11))') " y  dispersions -> ",disp(3,:)
  write(mf,'(a19,3(1x,D18.11))') " y' dispersions -> ",disp(4,:)

 
 do i=1,psr%n

  call track(psr,ray,i,i+1,state)

  a_t=ray
 
 do j=1,no1
  je=0
  je(5)=j
  do k=1,4
   disp(k,j)=a_t%v(k).sub.je  ! All the dispersions up to delta^3 are stored in disp(1:4,1:3)
 enddo
 enddo

  write(mf,*) " After Element ",f%mag%name
  write(mf,'(a19,3(1x,D18.11))') " x  dispersions -> ",disp(1,:)
  write(mf,'(a19,3(1x,D18.11))') " x' dispersions -> ",disp(2,:)
  write(mf,'(a19,3(1x,D18.11))') " y  dispersions -> ",disp(3,:)
  write(mf,'(a19,3(1x,D18.11))') " y' dispersions -> ",disp(4,:)

 f=>f%next
 enddo
 
close(mf)

write(6,*) "   "
write(6,*) " hit return to terminate program " 
write(6,*) "   "
pause 
call kill(id,m,A_T)  
call kill(ray)
call kill(normal)


end program program_psr


!=================================================================

subroutine  build_tiny_package_example(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut,H
type(fibre) ::  sf, L, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
call update_states

madlength = .false.
exact_model = .false.
madkind2 = drift_kick_drift

ang = (twopi / 20.d0)
Larc = 1.d0
brho = 1.2d0 * (Larc / ang)
h=  twopi/20/Larc
call set_mad(brho = brho, method = 2, step = 100)

kf = 1.d0
kd = -1.d0

L=DRIFT("L",Larc)
qf = sbend("QF", Larc, ang).q.kf  ! <-- PUTS QUADRUPOLE KF
qd = sbend("QD", Larc, ang).q.kD  ! <-- PUTS QUADRUPOLE KD

sf = sextupole("SF", 0.d0, 1.d0) 


cell = QF+SF+L+QD+SF+L

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_tiny_package_example

subroutine  build_PSR(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut
type(fibre) :: b, d1, d2, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
exact_model = .true.
default = default + nocavity  
call update_states
madlength = .false.

ang = (twopi * 36.d0 / 360.d0)
Larc = 2.54948d0
brho = 1.2d0 * (Larc / ang)
call set_mad(brho = brho, method = 2, step = 10)
madkind2 = drift_kick_drift

kf =  2.72d0 / brho
kd = -1.92d0 / brho

d1 = drift("D1", 2.28646d0)
d2 = drift("D2", 0.45d0)
qf = quadrupole("QF", 0.5d0, kf)
qd = quadrupole("QD", 0.5d0, kd)
b  = rbend("B", Larc, ang)
cell = d1 + qd + d2 + b + d2 + qf + d1

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_PSR
